#
from typing import Dict
import matplotlib.pyplot as plt
import pandas as pd
from model import Kronos, KronosTokenizer, KronosPredictor
from finetune.qlib_data_preprocess import QlibDataPreprocessor

class KronosEngine(object):
    def __init__(self):
        self.name = 'apps.kronos.kronos_engine.KronosEngine'

    @staticmethod
    def predict(params:Dict = {}) -> None:
        # 1. Load Model and Tokenizer
        tokenizer = KronosTokenizer.from_pretrained("NeoQuasar/Kronos-Tokenizer-base")
        model = Kronos.from_pretrained("NeoQuasar/Kronos-small")
        # 2. Instantiate Predictor
        predictor = KronosPredictor(model, tokenizer, device="cuda:0", max_context=512)
        # 3. Prepare Data
        df = pd.read_csv("anns/kronos/examples/data/XSHG_5min_600977.csv")
        df['timestamps'] = pd.to_datetime(df['timestamps'])
        lookback = 400
        pred_len = 120
        x_df = df.loc[:lookback-1, ['open', 'high', 'low', 'close', 'volume', 'amount']]
        x_timestamp = df.loc[:lookback-1, 'timestamps']
        y_timestamp = df.loc[lookback:lookback+pred_len-1, 'timestamps']
        # 4. Make Prediction
        pred_df = predictor.predict(
            df=x_df,
            x_timestamp=x_timestamp,
            y_timestamp=y_timestamp,
            pred_len=pred_len,
            T=1.0,
            top_p=0.9,
            sample_count=1,
            verbose=True
        )
        # 5. Visualize Results
        print("Forecasted Data Head:")
        print(pred_df.head())
        # Combine historical and forecasted data for plotting
        kline_df = df.loc[:lookback+pred_len-1]
        # visualize
        KronosEngine.plot_prediction(kline_df, pred_df)

    @staticmethod
    def plot_prediction(kline_df, pred_df):
        pred_df.index = kline_df.index[-pred_df.shape[0]:]
        sr_close = kline_df['close']
        sr_pred_close = pred_df['close']
        sr_close.name = 'Ground Truth'
        sr_pred_close.name = "Prediction"
        sr_volume = kline_df['volume']
        sr_pred_volume = pred_df['volume']
        sr_volume.name = 'Ground Truth'
        sr_pred_volume.name = "Prediction"
        close_df = pd.concat([sr_close, sr_pred_close], axis=1)
        volume_df = pd.concat([sr_volume, sr_pred_volume], axis=1)
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 6), sharex=True)
        ax1.plot(close_df['Ground Truth'], label='Ground Truth', color='blue', linewidth=1.5)
        ax1.plot(close_df['Prediction'], label='Prediction', color='red', linewidth=1.5)
        ax1.set_ylabel('Close Price', fontsize=14)
        ax1.legend(loc='lower left', fontsize=12)
        ax1.grid(True)
        ax2.plot(volume_df['Ground Truth'], label='Ground Truth', color='blue', linewidth=1.5)
        ax2.plot(volume_df['Prediction'], label='Prediction', color='red', linewidth=1.5)
        ax2.set_ylabel('Volume', fontsize=14)
        ax2.legend(loc='upper left', fontsize=12)
        ax2.grid(True)
        plt.tight_layout()
        plt.show()

    @staticmethod
    def finetune(params:Dict = {}) -> None:
        # This block allows the script to be run directly to perform data preprocessing.
        preprocessor = QlibDataPreprocessor()
        preprocessor.initialize_qlib()
        preprocessor.load_qlib_data()
        preprocessor.prepare_dataset()
